#!/bin/bash

# =========================================================
# Ubuntu 22.04 - Python + Playwright + Chromium + Xvfb 安装脚本（菜单版）
# 自动检测浏览器版本
# 支持 headless 和虚拟屏幕运行
# =========================================================

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 显示Banner
show_banner() {
    clear
    echo -e "${BLUE}=========================================================${NC}"
    echo -e "${GREEN}   Ubuntu Agent 安装脚本 - 交互式菜单版${NC}"
    echo -e "${BLUE}=========================================================${NC}"
    echo ""
}

# 步骤1: 更新系统
step1_update_system() {
    echo -e "${YELLOW}🚀 [1/10] 更新系统...${NC}"
    # ✅ 使用非交互式参数，自动保持本地配置文件版本（避免配置对话框）
    sudo DEBIAN_FRONTEND=noninteractive apt update -y
    sudo DEBIAN_FRONTEND=noninteractive apt upgrade -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold"
    echo -e "${GREEN}✅ 系统更新完成${NC}"
    read -p "按回车键继续..."
}

# 步骤2: 安装Python
step2_install_python() {
    echo -e "${YELLOW}📦 [2/10] 安装 Python 与基础依赖...${NC}"
    sudo DEBIAN_FRONTEND=noninteractive apt install -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" \
        python3 python3-pip python3-venv python-is-python3
    # ✅ 验证 pip 是否安装成功
    if ! command -v pip3 &> /dev/null && ! command -v pip &> /dev/null; then
        echo -e "${RED}❌ pip 安装失败，尝试手动修复...${NC}"
        sudo apt install --reinstall -y python3-pip || true
    fi
    echo -e "${GREEN}✅ 当前 Python 版本:${NC}"
    python --version
    if command -v pip &> /dev/null; then
        echo -e "${GREEN}✅ pip 版本: $(pip --version)${NC}"
    elif command -v pip3 &> /dev/null; then
        echo -e "${GREEN}✅ pip3 版本: $(pip3 --version)${NC}"
    fi
    read -p "按回车键继续..."
}

# 步骤3: 安装中文字体
step3_install_fonts() {
    echo -e "${YELLOW}📦 [3/10] 安装中文字体...${NC}"
    sudo apt-get install -y \
        fonts-wqy-zenhei \
        fonts-wqy-microhei \
        fonts-noto-cjk \
        fonts-liberation

    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✅ 中文字体安装成功${NC}"
        
        # 刷新字体缓存
        echo "🔄 刷新字体缓存..."
        fc-cache -fv > /dev/null 2>&1
        echo -e "${GREEN}✅ 字体缓存已更新${NC}"
        
        # 验证中文字体
        echo "🔍 验证中文字体安装..."
        if fc-list :lang=zh | grep -q "WenQuanYi\|Noto"; then
            echo -e "${GREEN}✅ 中文字体验证通过${NC}"
        else
            echo -e "${YELLOW}⚠️  中文字体可能未正确安装${NC}"
        fi
    else
        echo -e "${RED}⚠️  中文字体安装失败，页面可能显示乱码${NC}"
    fi
    read -p "按回车键继续..."
}

# 步骤4: 安装Playwright系统依赖
step4_install_playwright_deps() {
    echo -e "${YELLOW}🧩 [4/10] 安装 Playwright 所需系统依赖...${NC}"
    # ✅ 修复：Ubuntu 新版本中 libasound2 被替换为 libasound2t64
    # ✅ 修复：libatk1.0-0 和 libcups2 也使用新版本
    # ✅ 使用非交互式参数，避免配置对话框
    sudo DEBIAN_FRONTEND=noninteractive apt install -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" \
        wget git curl \
        libnss3 libx11-xcb1 libxcomposite1 libxdamage1 libxrandr2 \
        libxkbcommon0 libpango-1.0-0 libatk1.0-0t64 libcups2t64 \
        libgbm1 libasound2t64 fonts-liberation libgtk-3-0 \
        libdrm2 libxshmfence1 libxfixes3 libxext6 2>&1 | grep -v "Note, selecting" || true
    echo -e "${GREEN}✅ Playwright 系统依赖安装完成${NC}"
    read -p "按回车键继续..."
}

# 步骤5: 安装Xvfb
step5_install_xvfb() {
    echo -e "${YELLOW}🖥️ [5/10] 安装虚拟屏幕 Xvfb ...${NC}"
    sudo apt install -y xvfb
    echo -e "${GREEN}✅ Xvfb 安装完成${NC}"
    read -p "按回车键继续..."
}

# 步骤6: 创建浏览器缓存目录
step6_create_browser_dir() {
    echo -e "${YELLOW}📂 [6/10] 创建统一浏览器缓存目录 /opt/playwright-browsers ...${NC}"
    sudo mkdir -p /opt/playwright-browsers
    sudo chmod -R 777 /opt/playwright-browsers

    # 永久写入环境变量（只写一次）
    if ! grep -q "PLAYWRIGHT_BROWSERS_PATH" ~/.bashrc; then
        echo 'export PLAYWRIGHT_BROWSERS_PATH=/opt/playwright-browsers' >> ~/.bashrc
        echo -e "${GREEN}✅ 环境变量已添加到 ~/.bashrc${NC}"
    else
        echo -e "${GREEN}✅ 环境变量已存在${NC}"
    fi
    export PLAYWRIGHT_BROWSERS_PATH=/opt/playwright-browsers
    echo -e "${GREEN}✅ 浏览器缓存目录创建完成${NC}"
    read -p "按回车键继续..."
}

# 步骤7: 创建Python虚拟环境
step7_create_venv() {
    echo -e "${YELLOW}🌱 [7/10] 创建虚拟环境 /opt/envs/playenv ...${NC}"
    sudo mkdir -p /opt/envs && sudo chmod -R 777 /opt/envs
    if [ ! -d "/opt/envs/playenv" ]; then
        python -m venv /opt/envs/playenv
        echo -e "${GREEN}✅ 虚拟环境创建成功${NC}"
    else
        echo -e "${GREEN}✅ 虚拟环境已存在${NC}"
    fi
    read -p "按回车键继续..."
}

# 步骤8: 安装Playwright
step8_install_playwright() {
    echo -e "${YELLOW}📦 [8/10] 安装 Playwright...${NC}"
    
    # 激活虚拟环境
    if [ -f "/opt/envs/playenv/bin/activate" ]; then
        source /opt/envs/playenv/bin/activate
        echo -e "${GREEN}✅ 虚拟环境已激活${NC}"
    else
        echo -e "${RED}❌ 虚拟环境不存在，请先执行步骤7${NC}"
        read -p "按回车键继续..."
        return
    fi
    
    pip install -U pip setuptools
    pip install playwright loguru
    echo -e "${GREEN}✅ Playwright 安装完成${NC}"
    read -p "按回车键继续..."
}

# 步骤9: 安装Chromium浏览器
step9_install_chromium() {
    echo -e "${YELLOW}🌐 [9/10] 安装 Chromium 浏览器...${NC}"
    
    # 激活虚拟环境
    if [ -f "/opt/envs/playenv/bin/activate" ]; then
        source /opt/envs/playenv/bin/activate
    else
        echo -e "${RED}❌ 虚拟环境不存在，请先执行步骤7和步骤8${NC}"
        read -p "按回车键继续..."
        return
    fi
    
    export PLAYWRIGHT_BROWSERS_PATH=/opt/playwright-browsers
    PLAYWRIGHT_BROWSERS_PATH=/opt/playwright-browsers playwright install chromium

    # 🔍 自动检测 Chromium 版本号
    CHROMIUM_DIR=$(find /opt/playwright-browsers -maxdepth 1 -type d -name "chromium-*" | head -n 1)
    HEADLESS_DIR=$(find /opt/playwright-browsers -maxdepth 1 -type d -name "chromium_headless_shell-*" | head -n 1)

    if [ -z "$CHROMIUM_DIR" ]; then
        echo -e "${RED}❌ Chromium 安装失败，请检查网络连接或 Playwright 版本。${NC}"
        read -p "按回车键继续..."
        return
    fi

    echo -e "${GREEN}✅ Chromium 安装成功:${NC}"
    echo "   - 正常版: $CHROMIUM_DIR"
    echo "   - Headless 版: ${HEADLESS_DIR:-未检测到}"

    # 修正权限
    sudo chmod -R 777 /opt/playwright-browsers
    echo -e "${GREEN}✅ 浏览器权限设置完成${NC}"
    read -p "按回车键继续..."
}

# 步骤10: 测试运行
step10_test_run() {
    echo -e "${YELLOW}🧪 [10/10] 测试运行中（通过 Xvfb 支持非 headless 模式）...${NC}"
    
    # 激活虚拟环境
    if [ -f "/opt/envs/playenv/bin/activate" ]; then
        source /opt/envs/playenv/bin/activate
    else
        echo -e "${RED}❌ 虚拟环境不存在，请先完成前面的步骤${NC}"
        read -p "按回车键继续..."
        return
    fi
    
    export PLAYWRIGHT_BROWSERS_PATH=/opt/playwright-browsers
    
    xvfb-run -a python - <<'PYCODE'
import asyncio
from playwright.async_api import async_playwright

async def test():
    async with async_playwright() as p:
        # headless=False 使用虚拟屏幕
        browser = await p.chromium.launch(headless=False)
        page = await browser.new_page()
        await page.goto("https://www.baidu.com", timeout=20000)
        print("✅ 网页标题:", await page.title())
        await browser.close()

asyncio.run(test())
PYCODE

    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✅ 测试成功！${NC}"
    else
        echo -e "${RED}❌ 测试失败${NC}"
    fi
    read -p "按回车键继续..."
}

# 一键自动安装所有步骤
auto_install_all() {
    echo -e "${GREEN}======================================${NC}"
    echo -e "${GREEN}开始一键自动安装...${NC}"
    echo -e "${GREEN}======================================${NC}"
    echo ""
    
    # 设置为非交互模式
    set -e
    
    step1_update_system_auto
    step2_install_python_auto
    step3_install_fonts_auto
    step4_install_playwright_deps_auto
    step5_install_xvfb_auto
    step6_create_browser_dir_auto
    step7_create_venv_auto
    step8_install_playwright_auto
    step9_install_chromium_auto
    step10_test_run_auto
    
    echo ""
    echo -e "${GREEN}======================================${NC}"
    echo -e "${GREEN}🎉 一键安装完成！${NC}"
    echo -e "${GREEN}======================================${NC}"
    show_final_info
    read -p "按回车键返回主菜单..."
}

# 自动模式的步骤（不需要用户按回车）
step1_update_system_auto() {
    echo -e "${YELLOW}🚀 [1/10] 更新系统...${NC}"
    # ✅ 使用非交互式参数，自动保持本地配置文件版本（避免配置对话框）
    sudo DEBIAN_FRONTEND=noninteractive apt update -y
    sudo DEBIAN_FRONTEND=noninteractive apt upgrade -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold"
    echo -e "${GREEN}✅ 系统更新完成${NC}"
}

step2_install_python_auto() {
    echo -e "${YELLOW}📦 [2/10] 安装 Python 与基础依赖...${NC}"
    sudo DEBIAN_FRONTEND=noninteractive apt install -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" \
        python3 python3-pip python3-venv python-is-python3
    # ✅ 验证 pip 是否安装成功
    if ! command -v pip3 &> /dev/null && ! command -v pip &> /dev/null; then
        echo -e "${RED}❌ pip 安装失败，尝试手动修复...${NC}"
        sudo apt install --reinstall -y python3-pip || true
    fi
    # ✅ 验证 Python 版本
    if command -v python &> /dev/null; then
    echo -e "${GREEN}✅ Python 安装完成: $(python --version)${NC}"
    else
        echo -e "${RED}❌ Python 安装失败${NC}"
        exit 1
    fi
}

step3_install_fonts_auto() {
    echo -e "${YELLOW}📦 [3/10] 安装中文字体...${NC}"
    sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" \
        fonts-wqy-zenhei fonts-wqy-microhei fonts-noto-cjk fonts-liberation > /dev/null 2>&1
    fc-cache -fv > /dev/null 2>&1
    echo -e "${GREEN}✅ 中文字体安装完成${NC}"
}

step4_install_playwright_deps_auto() {
    echo -e "${YELLOW}🧩 [4/10] 安装 Playwright 所需系统依赖...${NC}"
    # ✅ 修复：Ubuntu 新版本中 libasound2 被替换为 libasound2t64
    # ✅ 修复：libatk1.0-0 和 libcups2 也使用新版本
    # ✅ 使用非交互式参数，避免配置对话框
    sudo DEBIAN_FRONTEND=noninteractive apt install -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" \
        wget git curl libnss3 libx11-xcb1 libxcomposite1 libxdamage1 libxrandr2 \
        libxkbcommon0 libpango-1.0-0 libatk1.0-0t64 libcups2t64 libgbm1 libasound2t64 fonts-liberation \
        libgtk-3-0 libdrm2 libxshmfence1 libxfixes3 libxext6 2>&1 | grep -v "Note, selecting" || true
    echo -e "${GREEN}✅ Playwright 系统依赖安装完成${NC}"
}

step5_install_xvfb_auto() {
    echo -e "${YELLOW}🖥️ [5/10] 安装虚拟屏幕 Xvfb ...${NC}"
    sudo DEBIAN_FRONTEND=noninteractive apt install -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" \
        xvfb > /dev/null 2>&1
    echo -e "${GREEN}✅ Xvfb 安装完成${NC}"
}

step6_create_browser_dir_auto() {
    echo -e "${YELLOW}📂 [6/10] 创建浏览器缓存目录...${NC}"
    sudo mkdir -p /opt/playwright-browsers
    sudo chmod -R 777 /opt/playwright-browsers
    if ! grep -q "PLAYWRIGHT_BROWSERS_PATH" ~/.bashrc; then
        echo 'export PLAYWRIGHT_BROWSERS_PATH=/opt/playwright-browsers' >> ~/.bashrc
    fi
    export PLAYWRIGHT_BROWSERS_PATH=/opt/playwright-browsers
    echo -e "${GREEN}✅ 浏览器缓存目录创建完成${NC}"
}

step7_create_venv_auto() {
    echo -e "${YELLOW}🌱 [7/10] 创建虚拟环境...${NC}"
    sudo mkdir -p /opt/envs && sudo chmod -R 777 /opt/envs
    if [ ! -d "/opt/envs/playenv" ]; then
        python -m venv /opt/envs/playenv
    fi
    echo -e "${GREEN}✅ 虚拟环境创建完成${NC}"
}

step8_install_playwright_auto() {
    echo -e "${YELLOW}📦 [8/10] 安装 Playwright...${NC}"
    
    # ✅ 检查虚拟环境是否存在
    if [ ! -f "/opt/envs/playenv/bin/activate" ]; then
        echo -e "${RED}❌ 虚拟环境不存在，请先执行步骤7${NC}"
        return 1
    fi
    
    # ✅ 激活虚拟环境
    source /opt/envs/playenv/bin/activate
    
    # ✅ 验证 pip 是否可用（虚拟环境中的 pip）
    if ! command -v pip &> /dev/null; then
        echo -e "${RED}❌ 虚拟环境中找不到 pip，尝试使用 python -m pip...${NC}"
        python -m pip install -U pip setuptools > /dev/null 2>&1
        python -m pip install playwright loguru > /dev/null 2>&1
    else
    pip install -U pip setuptools > /dev/null 2>&1
    pip install playwright loguru > /dev/null 2>&1
    fi
    
    echo -e "${GREEN}✅ Playwright 安装完成${NC}"
}

step9_install_chromium_auto() {
    echo -e "${YELLOW}🌐 [9/10] 安装 Chromium 浏览器...${NC}"
    source /opt/envs/playenv/bin/activate
    export PLAYWRIGHT_BROWSERS_PATH=/opt/playwright-browsers
    PLAYWRIGHT_BROWSERS_PATH=/opt/playwright-browsers playwright install chromium > /dev/null 2>&1
    sudo chmod -R 777 /opt/playwright-browsers
    echo -e "${GREEN}✅ Chromium 浏览器安装完成${NC}"
}

step10_test_run_auto() {
    echo -e "${YELLOW}🧪 [10/10] 测试运行...${NC}"
    source /opt/envs/playenv/bin/activate
    export PLAYWRIGHT_BROWSERS_PATH=/opt/playwright-browsers
    xvfb-run -a python - <<'PYCODE' > /dev/null 2>&1
import asyncio
from playwright.async_api import async_playwright

async def test():
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=False)
        page = await browser.new_page()
        await page.goto("https://www.baidu.com", timeout=20000)
        await browser.close()

asyncio.run(test())
PYCODE
    
    if [ $? -eq 0 ]; then
        echo -e "${GREEN}✅ 测试成功！${NC}"
    else
        echo -e "${RED}❌ 测试失败${NC}"
    fi
}

# 显示最终信息
show_final_info() {
    echo ""
    echo -e "${BLUE}=========================================================${NC}"
    echo -e "${GREEN}📍 虚拟环境: /opt/envs/playenv${NC}"
    echo -e "${GREEN}📍 浏览器路径: /opt/playwright-browsers${NC}"
    echo -e "${GREEN}📍 激活命令: source /opt/envs/playenv/bin/activate${NC}"
    echo -e "${GREEN}🖥️  运行脚本: xvfb-run -a python your_script.py${NC}"
    echo -e "${BLUE}=========================================================${NC}"
}

# 主菜单
show_menu() {
    show_banner
    echo -e "${YELLOW}请选择操作：${NC}"
    echo ""
    echo "  1)  🚀 一键自动安装（推荐）"
    echo "  ─────────────────────────────────────"
    echo "  2)  🚀 [1/10] 更新系统"
    echo "  3)  📦 [2/10] 安装 Python 与基础依赖"
    echo "  4)  📦 [3/10] 安装中文字体"
    echo "  5)  🧩 [4/10] 安装 Playwright 系统依赖"
    echo "  6)  🖥️  [5/10] 安装虚拟屏幕 Xvfb"
    echo "  7)  📂 [6/10] 创建浏览器缓存目录"
    echo "  8)  🌱 [7/10] 创建 Python 虚拟环境"
    echo "  9)  📦 [8/10] 安装 Playwright"
    echo "  10) 🌐 [9/10] 安装 Chromium 浏览器"
    echo "  11) 🧪 [10/10] 测试运行"
    echo "  ─────────────────────────────────────"
    echo "  12) ℹ️  显示安装信息"
    echo "  0)  ❌ 退出脚本"
    echo ""
    echo -ne "${GREEN}请输入选项 [0-12]: ${NC}"
}

# 主循环
main() {
    while true; do
        show_menu
        read -r choice
        
        case $choice in
            1)
                auto_install_all
                ;;
            2)
                step1_update_system
                ;;
            3)
                step2_install_python
                ;;
            4)
                step3_install_fonts
                ;;
            5)
                step4_install_playwright_deps
                ;;
            6)
                step5_install_xvfb
                ;;
            7)
                step6_create_browser_dir
                ;;
            8)
                step7_create_venv
                ;;
            9)
                step8_install_playwright
                ;;
            10)
                step9_install_chromium
                ;;
            11)
                step10_test_run
                ;;
            12)
                clear
                show_final_info
                read -p "按回车键返回主菜单..."
                ;;
            0)
                echo ""
                echo -e "${GREEN}感谢使用！再见！${NC}"
                echo ""
                exit 0
                ;;
            *)
                echo -e "${RED}无效选项，请重新选择${NC}"
                sleep 2
                ;;
        esac
    done
}

# 启动脚本
main

